import rclpy
from rclpy.node import Node
import numpy as np
import open3d as o3d
from sensor_msgs.msg import PointCloud2
from sensor_msgs_py import point_cloud2
import math
import time

class CableStatusAnalyzer(Node):
    def __init__(self):
        super().__init__('cable_status_analyzer')
        self.subscription = self.create_subscription(
            PointCloud2,
            'simulated_pointcloud',
            self.listener_callback,
            10)
        self.subscription
        
        # 可视化配置
        self.vis = o3d.visualization.VisualizerWithKeyCallback()
        self.vis.create_window(window_name="Cable Analyzer (Robust Mode)")
        opt = self.vis.get_render_option()
        opt.background_color = [0, 0, 0]
        opt.point_size = 5  # 增大点大小，便于观察小段钢缆
        
        # 数据存储
        self.clustered_pcd = o3d.geometry.PointCloud()
        self.coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
        self.vis.add_geometry(self.coord_frame)
        
        # 分析结果
        self.cable_params = {
            'left': {'length': 0.0, 'angle': 0.0, 'diameter': 0.0, 'detected': False},
            'right': {'length': 0.0, 'angle': 0.0, 'diameter': 0.0, 'detected': False}
        }
        self.first_update = True
        
        # 处理时间跟踪
        self.processing_time_ms = 0.0

        # 调整聚类参数以适应小段钢缆（恶劣环境优化）
        self.dbscan_eps = 0.4  # 增大邻域半径，适应稀疏点云
        self.dbscan_min_samples = 5  # 降低最小点数要求，适应小段钢缆
        self.min_cable_points = 50  # 大幅降低最小点数阈值，支持小段钢缆检测
        
        # 新增：钢缆检测置信度阈值
        self.confidence_threshold = 0.5

    def preprocess_points(self, points):
        """预处理点云：去噪和下采样，增强恶劣环境下的鲁棒性"""
        if len(points) < 50:
            return points
            
        # 使用统计滤波去除离群点
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        
        # 对稀疏点云使用更宽松的滤波参数
        cl, ind = pcd.remove_statistical_outlier(nb_neighbors=10, std_ratio=1.5)
        filtered_points = np.asarray(pcd.points)[ind]
        
        self.get_logger().info(f"预处理后点云: {len(filtered_points)}个点 (原始: {len(points)})")
        return filtered_points

    def cluster_cables_by_density(self, cable_points):
        """优化的密度聚类算法，适应小段钢缆检测"""
        if len(cable_points) < self.min_cable_points:
            self.get_logger().warn(f"钢缆点不足({len(cable_points)})")
            return None, None
        
        # 转换为Open3D点云
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(cable_points)
        
        # 密度聚类（使用更宽松的参数）
        with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Error):
            labels = np.array(pcd.cluster_dbscan(
                eps=self.dbscan_eps,
                min_points=self.dbscan_min_samples,
                print_progress=False
            ))
        
        # 提取有效聚类（排除噪声）
        unique_labels, counts = np.unique(labels, return_counts=True)
        valid_labels = unique_labels[unique_labels != -1]
        
        # 处理只有一个有效聚类的情况（可能只检测到一根钢缆）
        if len(valid_labels) == 1:
            self.get_logger().info(f"仅检测到1根钢缆")
            cluster1 = cable_points[labels == valid_labels[0]]
            return cluster1, np.array([])
        # 处理没有有效聚类的情况
        elif len(valid_labels) < 1:
            self.get_logger().warn(f"未检测到有效钢缆聚类，尝试备用方案")
            return self.fallback_split(cable_points)
        
        # 按点数排序，取前两个最大聚类
        sorted_indices = np.argsort(counts[unique_labels != -1])[::-1]
        main_labels = valid_labels[sorted_indices[:2]]
        
        cluster1 = cable_points[labels == main_labels[0]]
        cluster2 = cable_points[labels == main_labels[1]]
        
        # 标记左右钢缆
        return self.label_left_right(cluster1, cluster2)

    def fallback_split(self, cable_points):
        """增强的备用分割方案，适应各种恶劣情况"""
        # 直接使用x坐标中位数分割，不依赖z坐标（适应任意小段钢缆）
        x_median = np.median(cable_points[:, 0])
        left = cable_points[cable_points[:,0] < x_median]
        right = cable_points[cable_points[:,0] >= x_median]
        
        # 检查分割结果是否合理
        if len(left) < self.min_cable_points // 2:
            return np.array([]), right
        if len(right) < self.min_cable_points // 2:
            return left, np.array([])
            
        return left, right

    def label_left_right(self, cluster1, cluster2):
        """优化的左右钢缆标记方法，适应小段钢缆"""
        # 对于小段钢缆，直接使用所有点的x坐标中位数
        x1 = np.median(cluster1[:, 0]) if len(cluster1) > 0 else -np.inf
        x2 = np.median(cluster2[:, 0]) if len(cluster2) > 0 else np.inf
        
        if x1 < x2:
            return cluster1, cluster2
        else:
            return cluster2, cluster1

    def perform_clustering(self, points):
        """主聚类流程：移除货物依赖，专注于钢缆检测"""
        # 预处理点云（去噪）
        processed_points = self.preprocess_points(points)
        
        # 直接将所有点视为可能的钢缆点（不依赖货物检测）
        cable_points = processed_points
        
        if len(cable_points) < self.min_cable_points:
            return None, None
        
        left_cable, right_cable = self.cluster_cables_by_density(cable_points)
        
        # 过滤过小聚类，但保留更小的片段
        min_points = max(self.min_cable_points // 2, 10)  # 允许非常短的钢缆片段
        if len(left_cable) < min_points:
            left_cable = np.array([])
        if len(right_cable) < min_points:
            right_cable = np.array([])
        
        self.get_logger().info(
            f"聚类完成: 左钢缆={len(left_cable)}, 右钢缆={len(right_cable)}"
        )
        return left_cable, right_cable

    def fit_line(self, points):
        """优化的直线拟合，适应小段钢缆"""
        if len(points) < 10:  # 接受更少的点进行拟合
            return None, None, 0.0  # 增加置信度返回值
        
        mean = np.mean(points, axis=0)
        centered = points - mean
        u, s, vh = np.linalg.svd(centered)
        direction = vh[0]
        
        # 计算拟合置信度（基于奇异值）
        if s[0] == 0:
            confidence = 0.0
        else:
            confidence = 1.0 - (s[1] / s[0])  # 主方向越明显，置信度越高
        
        return direction, mean, confidence

    def analyze_cable(self, cable_points, is_left=True):
        """优化的钢缆分析，适应小段钢缆"""
        result = {'length': 0.0, 'angle': 0.0, 'diameter': 0.0, 'detected': False, 'confidence': 0.0}
        
        if len(cable_points) < self.min_cable_points // 2:
            return result
        
        direction, mean, confidence = self.fit_line(cable_points)
        if direction is None or confidence < self.confidence_threshold:
            return result
        
        # 修正方向向量（向下）
        if direction[2] < 0:
            direction = -direction
        
        # 长度计算（对小段钢缆更精确）
        projections = np.dot(cable_points - mean, direction)
        length = np.max(projections) - np.min(projections)
        
        # 偏角计算
        vertical_dir = np.array([0, 0, 1])
        dot_product = np.dot(direction, vertical_dir)
        dot_product = np.clip(dot_product, -1.0, 1.0)
        angle_deg = np.degrees(np.arccos(dot_product))
        
        # 直径计算（对小段钢缆使用稳健估计）
        centered = cable_points - mean
        perpendicular = centered - np.outer(projections, direction)
        distances = np.linalg.norm(perpendicular, axis=1)
        
        # 使用中位数代替平均值，减少异常值影响
        diameter = 2 * np.median(distances) if len(distances) > 0 else 0.0
        
        return {
            'length': length, 
            'angle': angle_deg, 
            'diameter': diameter,
            'detected': True,
            'confidence': confidence
        }

    def visualize_results(self, left, right):
        """可视化结果，适应可能只有部分钢缆的情况"""
        all_points = []
        all_colors = []
        
        # 左钢缆：红色
        if len(left) > 0:
            all_points.append(left)
            all_colors.append(np.tile([1, 0, 0], (len(left), 1)))
        
        # 右钢缆：绿色
        if len(right) > 0:
            all_points.append(right)
            all_colors.append(np.tile([0, 1, 0], (len(right), 1)))
        
        # 无钢缆时显示提示
        if not all_points:
            self.get_logger().warn("未检测到任何钢缆点")
            return
        
        all_points = np.vstack(all_points)
        all_colors = np.vstack(all_colors)
        
        self.clustered_pcd.points = o3d.utility.Vector3dVector(all_points)
        self.clustered_pcd.colors = o3d.utility.Vector3dVector(all_colors)
        
        if self.first_update:
            self.vis.add_geometry(self.clustered_pcd)
            self.first_update = False
        else:
            self.vis.update_geometry(self.clustered_pcd)
        
        self.vis.poll_events()
        self.vis.update_renderer()

    def display_analysis(self):
        """显示分析结果，包括置信度信息"""
        self.get_logger().info("\n===== 钢缆状态分析 =====")
        
        # 左钢缆信息
        left = self.cable_params['left']
        if left['detected']:
            self.get_logger().info(
                f"左钢缆: 长度={left['length']:.2f}m, "
                f"偏角={left['angle']:.2f}°, "
                f"直径={left['diameter']:.4f}m, "
                f"置信度={left['confidence']:.2f}"
            )
        else:
            self.get_logger().info("左钢缆: 未检测到")
        
        # 右钢缆信息
        right = self.cable_params['right']
        if right['detected']:
            self.get_logger().info(
                f"右钢缆: 长度={right['length']:.2f}m, "
                f"偏角={right['angle']:.2f}°, "
                f"直径={right['diameter']:.4f}m, "
                f"置信度={right['confidence']:.2f}"
            )
        else:
            self.get_logger().info("右钢缆: 未检测到")
        
        self.get_logger().info(f"处理时间: {self.processing_time_ms:.2f} ms")
        self.get_logger().info("=======================\n")

    def listener_callback(self, msg):
        """主回调函数，优化恶劣环境适应性"""
        start_time = time.time()
        
        # 读取点云数据
        points = point_cloud2.read_points_list(msg, field_names=("x", "y", "z"), skip_nans=True)
        point_array = np.array([(p.x, p.y, p.z) for p in points], dtype=np.float32)
        self.get_logger().info(f"接收点云: {len(point_array)}个点")
        
        # 即使点云很少也尝试分析（降低阈值）
        if len(point_array) < 30:
            self.get_logger().warn("点云数量极少，尝试分析...")
            # 仍然尝试分析，而不是直接返回
        
        # 执行聚类（专注于钢缆）
        left_cable, right_cable = self.perform_clustering(point_array)
        if left_cable is None:
            # 重置分析结果
            self.cable_params['left']['detected'] = False
            self.cable_params['right']['detected'] = False
            return
        
        # 分析钢缆参数
        self.cable_params['left'] = self.analyze_cable(left_cable)
        self.cable_params['right'] = self.analyze_cable(right_cable)
        
        # 可视化结果
        self.visualize_results(left_cable, right_cable)
        
        # 计算处理时间
        end_time = time.time()
        self.processing_time_ms = (end_time - start_time) * 1000
        
        # 显示分析结果
        self.display_analysis()

def main(args=None):
    rclpy.init(args=args)
    analyzer = CableStatusAnalyzer()
    try:
        rclpy.spin(analyzer)
    except KeyboardInterrupt:
        analyzer.vis.destroy_window()
    analyzer.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
